# Copyright (c) 2019 Presto Labs Pte. Ltd.
# Author: inkyu

import pandas as pd
import numpy as np
import scipy.stats
import datetime
import time
from io import BytesIO

import functools
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt

from coin.strategy.mm.fastfeature.ridge_regress import get_quantile, _chunks


def interval_plt(x, y, *, fmt, maxpoints=None, **kwargs):
  if maxpoints is not None and len(x) >= maxpoints * 2:
    every_n = max(int(len(x) / maxpoints), 1)
    x = x[::every_n]
    y = y[::every_n]
  plt.plot(x, y, fmt, **kwargs)


def timeseries_plt(t, a, **kwargs):
  df = pd.DataFrame({'a': a, 't': pd.DatetimeIndex(t)}).dropna()
  interval_plt(pd.DatetimeIndex(df['t']), df['a'], **kwargs)
  plt.xticks(rotation=20)


def x_y_scatter(x, y, **kwargs):
  df = pd.DataFrame({'x': x, 'y': y}).dropna()
  interval_plt(df['x'], df['y'], fmt='C2.', markersize=0.5, **kwargs)


x_timeseries_scatter = functools.partial(timeseries_plt, fmt='C1.', markersize=0.5)
y_timeseries_scatter = functools.partial(timeseries_plt, fmt='C0.', markersize=0.5)


def ret_per_trade_plot(x, y, *args, **kwargs):
  halfidx = int(len(x) / 2)
  xin = range(halfidx)
  xout = range(halfidx, len(x))

  qtl_list = [0.4, 0.3, 0.2, 0.1]
  color_list = ['C1', 'C3', 'C5', 'C7']
  legsdict = dict()

  for legcol, xidx, xdiff, yval, markeradd in [('is', xin, x[:halfidx], y[:halfidx], ':'),
                                               ('oos', xout, x[halfidx:], y[halfidx:], '-')]:
    xdiff = pd.Series(xdiff)
    for qtl, basemarker in zip(qtl_list, color_list):
      marker = basemarker + markeradd
      xsignal = ((xdiff < xdiff.quantile(qtl)) * (xdiff - xdiff.quantile(qtl)) +
                 (xdiff > xdiff.quantile(1 - qtl)) * (xdiff - xdiff.quantile(1 - qtl)))
      xsignal.iloc[0] = 0
      xsignal.iloc[-1] = 0

      turnover = xsignal.diff().abs().sum()
      ret = (xsignal * yval).sum()
      interval_plt(xidx, (xsignal * yval).cumsum(), fmt=marker, linewidth=0.5)
      if qtl not in legsdict:
        legsdict[qtl] = {}
      legsdict[qtl][legcol] = ret / turnover
  legs = [
      f'qtl:{key:.3f}|isret/t{row["is"]:.1f}|os{row["oos"]:.1f}' for key, row in legsdict.items()
  ]
  plt.legend(legs)


def plotting_init():
  plt.rcParams['lines.linewidth'] = 0.5
  plt.rcParams['figure.figsize'] = 15, 13
  plt.rcParams['font.size'] = 10
  plt.rcParams['legend.fontsize'] = 10
  plt.rcParams['xtick.labelsize'] = 10
  plt.rcParams['ytick.labelsize'] = 10


def plotting(x, y, t, p, title, desc):
  df = pd.DataFrame({'x': x, 'y': y, 't': t, 'p': p}).dropna()

  pnl = (df['x'] / df['p'] * df['y']).cumsum()
  turnover = df['x'].diff().abs().sum()
  df['date'] = pd.to_datetime(df['t']).dt.date
  days = df['date'].nunique()

  intv = int(len(pnl) / days / 10)
  sep_pnl = pnl.diff(periods=intv).dropna()[::intv]

  sharpe_intv = sep_pnl.mean() / sep_pnl.std()
  sharpe_str = "Sharpe: %f" % (sharpe_intv / np.sqrt(len(sep_pnl)) * np.sqrt(days))
  pnl_str = "PnL: %f" % pnl.iloc[-1]
  turnover_str = "Turnover: %f" % (turnover / days)

  desc += '\n' + ', '.join([pnl_str, sharpe_str, turnover_str])
  plt.subplots_adjust(left=0.05, hspace=0.3, wspace=0.25)
  plt.suptitle(title, fontsize=18, y=0.98)
  plt.figtext(.5, 0.90, f'Alpha Description\n{desc}', fontsize=12, ha='center')

  maxpoints = None

  plt.subplot(421)
  x_timeseries_scatter(df['t'], df['x'], maxpoints=maxpoints)
  plt.subplot(422)
  x_y_scatter(df['x'], df['y'], maxpoints=maxpoints)
  plt.subplot(423)
  y_timeseries_scatter(df['t'], df['y'], maxpoints=maxpoints)
  plt.subplot(424)
  x_y_quantile(df['x'], df['y'], df['t'])
  plt.subplot(425)
  qtl_plot(df['x'], df['y'], df['t'])
  plt.subplot(426)
  cum_qtl_plot(df['x'], df['y'])
  plt.subplot(427)

  timeseries_plt(df['t'], pnl, fmt='C4-', maxpoints=maxpoints, label='pnl')
  # plt.plot(pd.DatetimeIndex(df['t']), pnl, 'C4-', linewidth=1, label='pnl')
  plt.xticks(rotation=20)
  # loc='best' consumes time
  plt.legend(loc='upper left')
  plt.twinx()

  timeseries_plt(df['t'],
                 df['p'],
                 fmt='C5:',
                 maxpoints=maxpoints,
                 linewidth=1,
                 markersize=1,
                 label='mid_p')
  plt.legend(loc='upper left')
  plt.subplot(428)
  ret_per_trade_plot(df['x'], df['y'])

  img = BytesIO()
  plt.savefig(img, format='png')
  plt.close()
  img.seek(0)

  # img_file_path = f'/static/{title}_{datetime.date.today().strftime("%Y%m%d")}_{str(time.time())[-6:]}.png'
  # plt.savefig(f'coin/strategy/interval{img_file_path}')

  return img  # img_file_path


def print_matplotlib(x, y, t, p, title, desc=''):
  return plotting(x, y, t, p, title, desc)


def print_matplotlib_multiple(data, base_df, comp_df, title):
  plt.switch_backend('agg')
  plt.rcParams['lines.linewidth'] = 0.5
  plt.rcParams['figure.figsize'] = 15, 1 + (len(data) * 8)
  plt.rcParams['font.size'] = 10
  plt.rcParams['legend.fontsize'] = 10
  plt.rcParams['xtick.labelsize'] = 10
  plt.rcParams['ytick.labelsize'] = 10

  plt.subplots_adjust(left=0.05, hspace=0.3, wspace=0.25)
  plt.suptitle(title, fontsize=18, y=0.98)
  plt.figtext(.5, 0.90, f'Alpha Description\n{title}', fontsize=12, ha='center')

  rows = len(data) * 4
  columns = 2
  offset = 0

  for datum in data:
    name, x, y, t, desc = datum

    df = pd.DataFrame({'x': x, 'y': y, 't': t}).dropna()
    xq, yq = get_quantile(df['x'], df['y'])

    plt.subplot(rows, columns, offset + 1)
    x_timeseries_scatter(t, x)
    plt.subplot(rows, columns, offset + 2)
    x_y_scatter(x, y)
    plt.subplot(rows, columns, offset + 3)
    y_timeseries_scatter(t, y)
    plt.subplot(rows, columns, offset + 4)
    x_y_quantile(x, y, t)
    plt.subplot(rows, columns, offset + 5)
    qtl_plot(x, y, t)
    plt.subplot(rows, columns, offset + 6)
    cum_qtl_plot(x, y)
    plt.subplot(rows, columns, offset + 7)
    plt.plot(pd.DatetimeIndex(df['t']), (scipy.stats.zscore(df['x']) * df['y']).cumsum(),
             'C4-',
             linewidth=1,
             label='pnl')
    plt.legend(loc='best')
    plt.twinx()
    plt.plot(pd.DatetimeIndex(base_df['global_timestamp']),
             base_df['t.vwap.1m'],
             'C5:',
             linewidth=1,
             markersize=1,
             label='vwap_1m')
    plt.legend(loc='best')
    plt.subplot(rows, columns, offset + 8)

    ret_per_trade_plot(x, y)
    offset += 8

  img_file_path = f'static/{title}_{datetime.date.today().strftime("%Y%m%d")}_{str(time.time())[-6:]}.png'
  plt.savefig(f'coin/strategy/interval/{img_file_path}')
  plt.close()
  return img_file_path


def x_y_quantile(x, y, t):
  halfidx = int(len(x) / 2)
  dfs = []
  for x_domain, y_domain, marker in [(x[:halfidx], y[:halfidx], 'C3.-'),
                                     (x[halfidx:], y[halfidx:], 'C5.-')]:
    avg_xs, avg_ys = get_quantile(x_domain, y_domain)
    plt.plot(avg_xs, avg_ys, marker, markersize=5, linewidth=1)
  plt.legend(['qqplot front half', 'qqplot back half'], loc='upper left')


def qtl_plot(x, y, t):
  dfs = []
  qtl_list = [0.4, 0.3, 0.2, 0.1]
  for idx, qtl in enumerate(qtl_list):
    xvec = pd.Series(x)
    xsignal = ((xvec < xvec.quantile(qtl)) * (xvec - xvec.quantile(qtl)) +
               (xvec > xvec.quantile(1 - qtl)) * (xvec - xvec.quantile(1 - qtl)))
    xsignal.iloc[0] = 0
    xsignal.iloc[-1] = 0
    maxpos = max(np.max(xsignal), -np.min(xsignal))
    xsignal /= maxpos
    cumpnl = (xsignal * y).cumsum()
    plt.plot(pd.DatetimeIndex(t), cumpnl, f'C{idx}-')
  plt.legend([f"qtl {qtl}" for qtl in qtl_list], loc='upper left')
  plt.xticks(rotation=20)


def cum_qtl_plot(x, y):
  halfidx = int(len(x) / 2)
  nchunks = 30
  for x_domain, y_domain, marker in [(x[:halfidx], y[:halfidx], 'C3.-'),
                                     (x[halfidx:], y[halfidx:], 'C5.-')]:
    x, _, y = zip(*sorted(zip(x_domain, range(len(x_domain)), y_domain)))
    x = np.array(x)
    y = np.array(y)
    avg_xs = []
    avg_ys = []
    idxp = (x >= 0)
    for neg, subx, suby in [(True, x[~idxp], y[~idxp]), (False, x[idxp], y[idxp])]:
      sel_avg_ys = [np.mean(y_chunk) for y_chunk in _chunks(suby, nchunks)]
      sel_avg_xs = [np.mean(x_chunk) for x_chunk in _chunks(subx, nchunks)]
      avg_xs.extend(sel_avg_xs)
      if neg:
        avg_ys.extend(np.array(sel_avg_ys)[::-1].cumsum()[::-1])
      else:
        avg_ys.extend(np.array(sel_avg_ys).cumsum())
    plt.plot(avg_xs, avg_ys, marker, linewidth=1.2)
  plt.legend(['qqplot front half', 'qqplot back half'], loc='upper left')
